﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using TorchSharp;
using static TorchSharp.torch;

namespace NgpSharp
{
    public class Activation : torch.autograd.SingleTensorFunction<Activation>
    {
        public override string Name => "activation";

        public override torch.Tensor forward(torch.autograd.AutogradContext ctx, params object[] vars)
        {
            using var c = torch.NewDisposeScope();
            torch.Tensor x = (torch.Tensor)vars[0];
            ctx.save_for_backward([x.MoveToOuterDisposeScope()]);
            var result = torch.exp(x);
            return result.MoveToOuterDisposeScope();
        }
        public override List<torch.Tensor> backward(torch.autograd.AutogradContext ctx, torch.Tensor g)
        {
            using var c = torch.NewDisposeScope();
            using torch.Tensor x = ctx.get_saved_variables()[0];
            var result = g * torch.exp(x.clamp(-15, 15));
            g.Dispose();
            return [result.MoveToOuterDisposeScope()];
        }
    }
}
